dumpling/export/writer_serial_test.go (2)

449-494: TestWriteInsertWithStatementSizeLimit and TestWriteInsertMultipleStatements are largely duplicative.

Both construct a mock table, set a small StatementSize, and assert (a) Count("INSERT INTO ... VALUES") > 1 and (b) each split chunk ends with ;. The only meaningful difference is the dataset/size tuple (6 rows @ 150 bytes vs 4 rows @ 50 bytes). Consider a single table-driven test to reduce maintenance cost and keep the suite minimal per the test guideline.

♻️ Sketch

func TestWriteInsertStatementSizeChunking(t *testing.T) {
    cases := []struct {
        name          string
        data          [][]driver.Value
        colTypes      []string
        tableName     string
        statementSize uint64
    }{
        {"users_6rows_150", /*...*/, 150},
        {"items_4rows_50",  /*...*/, 50},
    }
    for _, tc := range cases {
        t.Run(tc.name, func(t *testing.T) { /* shared assertions */ })
    }
}

As per coding guidelines: "Keep test changes minimal and deterministic; avoid broad golden/testdata churn unless required."

Also applies to: 535-571

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/writer_serial_test.go` around lines 449 - 494, Two tests
(TestWriteInsertWithStatementSizeLimit and TestWriteInsertMultipleStatements)
duplicate logic by creating mock tableIR data, setting a small statementSize,
calling WriteInsert and asserting multiple INSERT chunks and semicolon
termination; consolidate them into a single table-driven test (e.g.,
TestWriteInsertStatementSizeChunking) that iterates cases with different
data/statementSize tuples, reuses the same setup (createMockConfig,
newMockTableIR, NewBufferWriter, configForWriteSQL, WriteInsert, and metrics)
and performs the shared assertions (count of "INSERT INTO `...` VALUES", each
chunk ends with ';', and row count) inside t.Run subtests to remove duplication
and keep tests deterministic.

---

449-494: Harden the row-presence assertion.

Line 493 counts ( characters as a proxy for row count. It happens to work because none of the test values contain (, but it will silently misreport if test data ever includes a parenthesis, and it does not actually prove that every input row made it to the output across chunk boundaries. Prefer using the returned n (already verified) plus a per-chunk row count, e.g. splitting on the INSERT prefix, trimming the ;\n, and summing the number of ),\n/);\n row terminators. At minimum, count ),\n + );\n instead of (.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/writer_serial_test.go` around lines 449 - 494, The current
row-presence assertion in TestWriteInsertWithStatementSizeLimit is fragile
because it counts "(" characters; instead, verify rows by parsing the generated
INSERT chunks: split output on the INSERT prefix ("INSERT INTO `users` VALUES"),
for each non-empty chunk count row terminators by summing occurrences of "),\n"
and ");\n" (or trimming the trailing ";\n" and counting occurrences of "),\n"
plus one per non-empty chunk if needed), and assert that this total equals the
returned n (or 6); update the assertion that currently uses
strings.Count(output, "(") to use this more robust per-chunk terminator counting
so WriteInsert and TestWriteInsertWithStatementSizeLimit are resilient to
parentheses in data and chunk boundaries.

dumpling/export/sql_util.go (2)

124-129: Stale godoc fragment above escapeSQLString.

The // pickupPossibleFieldsForStringChunking returns all columns… line documents a different (removed?) function but is physically attached to escapeSQLString, which confuses go doc output and readers. Drop that line (or move it if the function still exists elsewhere).

♻️ Proposed cleanup

-// pickupPossibleFieldsForStringChunking returns all columns of the selected index for composite key chunking
-// escapeSQLString properly escapes a string for use in internal SQL boundary condition queries.
+// escapeSQLString properly escapes a string for use in internal SQL boundary condition queries.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/sql_util.go` around lines 124 - 129, The godoc above
escapeSQLString contains a stale fragment referencing
pickupPossibleFieldsForStringChunking; remove that stray line (or relocate it to
the correct function if that function still exists) so the comment block
immediately preceding escapeSQLString documents only escapeSQLString; update the
comment text to accurately describe escapeSQLString and ensure no leftover
references to pickupPossibleFieldsForStringChunking remain.

---

94-116: Unique-key tie-break is non-deterministic.

When multiple unique keys share the same column count, the chosen bestKey depends on Go map iteration order, so the chunking key can differ run-to-run for the same schema. If determinism matters for reproducibility/debugging of chunk boundaries, break ties by keyName (lexicographic) or by the first key seen in SHOW INDEX order.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/sql_util.go` around lines 94 - 116, The current tie-break for
selecting bestKey from uniqueKeyMap is non-deterministic because map iteration
order can vary; update the selection logic in the block that iterates
uniqueKeyMap so ties on keyInfo.count are broken deterministically: collect the
map keys, sort them lexicographically (keyName), then iterate the sorted key
list to pick the key with the smallest count (if equal counts, the
lexicographically smallest keyName wins). Ensure you update references to
uniqueKeyMap, bestKey and keyInfo.count so the final chosen bestKey is
deterministic across runs; alternatively sort by the original SHOW INDEX order
if that ordering is available.

dumpling/export/ir_impl_test.go (1)

149-185: Test name overstates coverage.

TestRowIterWithStringKeyProgress exercises newRowIter / Decode / HasNext / Close against sqlmock rows, which is fine, but there is nothing "string-key" or "progress"-specific about it — the behavior under test is identical for numeric rows, and progress counters are not touched. Consider renaming (e.g. TestRowIterStringKeyDecode) so failures point to the actual surface under test.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/ir_impl_test.go` around lines 149 - 185, Rename the test
function to accurately reflect what's being tested: change
TestRowIterWithStringKeyProgress to a name like TestRowIterStringKeyDecode (or
TestRowIterDecode) to indicate it exercises newRowIter, Decode, HasNext and
Close against sqlmock rows; update the test function declaration and any
references to that identifier so failures point to the correct surface under
test.

dumpling/export/dump_test.go (1)

1047-1083: Optional: add compile-time check that mockTableMeta satisfies TableMeta.

The comment claims the mock implements TableMeta, but only a handful of methods are defined and nothing enforces the claim. If the real interface grows a method, this will not break and silently drift. A one-liner var _ TableMeta = (*mockTableMeta)(nil) makes the contract explicit; alternatively drop the comment since tests only use ChunkKey().

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/dump_test.go` around lines 1047 - 1083, Add a compile-time
assertion that mockTableMeta implements the TableMeta interface by adding a var
declaration like var _ TableMeta = (*mockTableMeta)(nil) near the mockTableMeta
type; this ensures the compiler will fail if TableMeta gains new methods and
keeps the test mock in sync with the real interface (reference mockTableMeta and
TableMeta).

dumpling/export/string_chunking_test.go (1)

107-130: Consider asserting the empty-input behavior matches intent.

{"", []string{""}} locks in that extractOrderByColumns("") returns a one-element slice containing "" rather than an empty slice. That is fine as a contract, but callers that do len(cols) > 0 to decide whether an ORDER BY exists will incorrectly see "one column". A short comment next to this case, or a helper that returns nil on empty input, would prevent the mis-use.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/string_chunking_test.go` around lines 107 - 130, The test
currently asserts that extractOrderByColumns("") returns a one-element slice
[""], which can mislead callers that check len(cols)>0; update
TestExtractOrderByColumns to either expect an empty (nil) slice for the
empty-input case (replace {"", []string{""}} with {"", []string{}} or nil
expectation) and add a one-line comment in the test beside that case clarifying
the contract for empty input, or alternatively change extractOrderByColumns to
return nil on empty input and update the test accordingly; reference the test
TestExtractOrderByColumns and the function extractOrderByColumns when making the
change.

dumpling/export/dump.go (1)

2161-2166: Stale comment in newTaskTableData.

// Chunking mode is already set at table level in concurrentDumpTable

Nothing visible in newTaskTableData or NewTaskTableData consults a "chunking mode" flag, and concurrentDumpTable doesn't set one. Looks like a leftover from an earlier iteration — please either remove it or update it to describe what the wrapper actually adds (the totalChunks.Add(1) bookkeeping).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/dump.go` around lines 2161 - 2166, The inline comment in
newTaskTableData is stale and misleading: remove or replace "// Chunking mode is
already set at table level in concurrentDumpTable" with a concise comment that
describes what this wrapper actually does (increments d.metrics.totalChunks and
constructs a TaskTableData via NewTaskTableData); ensure no other code relies on
a non-existent "chunking mode" flag in concurrentDumpTable or NewTaskTableData.

dumpling/export/string_chunking.go (1)

201-216: Sampling error downgrades silently; consider surfacing it.

On QuerySQL failure the loop just logs a Warn and breaks, then the post-loop branch (line 257+) appends a "rest of the table" tail chunk. That means the dump itself still succeeds but degrades to a single huge tail chunk — defeating the parallelism the feature exists to provide, without any error signal to the caller. Consider returning the error (or at least an AddCounter to a dedicated failure metric) so operators notice. If the current behavior is intentional (e.g., to stay compatible with best-effort master behavior), a short comment here explaining that contract would help future readers.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/string_chunking.go` around lines 201 - 216, The sampling loop
currently swallows QuerySQL errors (logging a Warn and break) which causes a
silent fallback to a single tail chunk; update the loop in string_chunking.go to
surface sampling failures: on QuerySQL error (the branch that logs via
tctx.L().Warn with db/tbl/i and err) either return the error up to the caller
(propagate the error from the function that runs the loop) or increment a
dedicated failure metric (e.g., tctx.AddCounter or a specific telemetry counter)
before breaking, and add a short comment stating the chosen contract
(best-effort vs. fail-fast); specifically change the error-handling where
currentBoundary and QuerySQL are used so callers/operators are notified instead
of silently degrading parallelism.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

dumpling/export/dump.go (2)

936-945: Refactor: reuse beginChunkTracking to remove the duplicated handshake.

The newly-introduced beginChunkTracking (lines 1019–1030) implements exactly the Store + finalize defer that this block does inline. Switching to the helper keeps the integer and TiDB‑sample paths in sync if the handshake ever changes.

♻️ Proposed refactor

 	chunkIndex := 0
 	nullValueCondition := ""
-	chunkStats := newTableChunkStat()
-	d.chunkedTables.Store(meta.ChunkKey(), chunkStats)
-	defer func() {
-		chunkStats.finalized.Store(true)
-		if chunkStats.finished.Load() == chunkStats.sent.Load() {
-			if _, loaded := d.chunkedTables.LoadAndDelete(meta.ChunkKey()); loaded {
-				IncCounter(d.metrics.finishedTablesCounter)
-			}
-		}
-	}()
+	defer d.beginChunkTracking(meta)()

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/dump.go` around lines 936 - 945, Replace the inline
chunk-tracking setup in dump.go (the new chunkStats := newTableChunkStat();
d.chunkedTables.Store(meta.ChunkKey(), chunkStats); defer { ... } block) with a
call to the existing helper beginChunkTracking so both paths share the same
Store + finalize defer logic; specifically, remove the manual
creation/storage/finalize defer and invoke beginChunkTracking(meta.ChunkKey(),
chunkStats) (or adapt to beginChunkTracking’s parameter list) so that chunkStats
and the finalize behavior are handled by beginChunkTracking, keeping integer and
TiDB-sample paths consistent.

---

883-896: Minor: log the fallback COUNT(*) failure at debug.

If QuerySQL fails or returns an invalid NullInt64, count is left at the EXPLAIN value and the function silently drops into the count < conf.Rows sequential branch — exactly the case this fallback was added to avoid. Logging the discarded error makes the silent downgrade diagnosable. As per coding guidelines: Keep error handling actionable and contextual; avoid silently swallowing errors.

📝 Proposed diff

-		err := conn.QuerySQL(tctx, func(rows *sql.Rows) error {
+		qErr := conn.QuerySQL(tctx, func(rows *sql.Rows) error {
 			return rows.Scan(&directCount)
 		}, func() { directCount = sql.NullInt64{} }, countQuery)
-		if err == nil && directCount.Valid && directCount.Int64 > 0 {
+		if qErr == nil && directCount.Valid && directCount.Int64 > 0 {
 			count = uint64(directCount.Int64)
+		} else if qErr != nil {
+			tctx.L().Debug("authoritative COUNT(*) failed; keeping EXPLAIN estimate",
+				zap.String("database", db), zap.String("table", tbl),
+				log.ShortError(qErr))
 		}

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/dump.go` around lines 883 - 896, When the fallback COUNT(*)
query fails or returns an invalid value, add a debug-level log so the downgrade
from EXPLAIN-derived count to sequential scan is observable; specifically, in
the block that calls QuerySQL (referencing QuerySQL, countQuery and directCount)
emit a debug log including the error when err != nil and also log when
directCount.Valid is false (mentioning the countQuery and conf.Rows/context),
then proceed with existing fallback behavior.

dumpling/export/string_chunking_test.go (1)

153-160: Optional: this test is tautological.

TestMaxChunkLimitSafety only re-asserts the literal value of the maxChunkLimit constant. It doesn't exercise the loop guard in streamStringChunks (line 117), so it gives no coverage of the safety behavior — only locks the constant in place. Either drop it, or flip it to assert that streamStringChunks actually breaks out of the boundary-sampling loop and emits a warning when the limit is reached (e.g., via a fake conn that always returns a new boundary).

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/string_chunking_test.go` around lines 153 - 160, The test
TestMaxChunkLimitSafety is tautological because it only asserts the literal
maxChunkLimit constant; update it to exercise the loop guard in
streamStringChunks instead: create a fake connection/mock that simulates an
infinite/new boundary scenario (always returns a new boundary or no progress)
and call streamStringChunks, then assert that the function breaks out once
maxChunkLimit is reached and emits the expected warning/log or closes the
stream; alternatively, if you prefer not to test runtime behavior, remove
TestMaxChunkLimitSafety to avoid locking the constant value in tests.

dumpling/export/string_chunking.go (1)

89-93: Minor: defer ignores sendBufferedChunk(true) error.

If the context is already cancelled when the defer runs, sendBufferedChunk(true) will return tctx.Err() and the buffered final chunk silently disappears. The Dump is shutting down anyway, but the dropped error makes the abort harder to diagnose. Either log it or rely on the surrounding tctx-cancellation path to surface it.

📝 Suggested change

-		// Send any remaining buffered chunk as the last chunk
-		if buffer != nil {
-			_ = sendBufferedChunk(true) // Mark as last chunk
-		}
+		// Send any remaining buffered chunk as the last chunk
+		if buffer != nil {
+			if err := sendBufferedChunk(true); err != nil {
+				tctx.L().Debug("failed to flush final string chunk on shutdown",
+					zap.String("database", db),
+					zap.String("table", tbl),
+					zap.Error(err))
+			}
+		}

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/string_chunking.go` around lines 89 - 93, Defer currently
calls sendBufferedChunk(true) and drops its error; change the defer to capture
the returned error and surface it (e.g., log it) unless it's the expected
context-cancellation error from tctx.Err(); specifically, in the deferred
closure around buffer and sendBufferedChunk, assign err :=
sendBufferedChunk(true) and if err != nil && err != tctx.Err() (or simply if err
!= nil) write a diagnostic via the package logger or fmt to record the failure
so the final chunk loss is visible for debugging.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

dumpling/export/string_chunking_test.go (1)

155-162: Redundant bounds assertions after exact equality.

Once require.Equal(t, int64(1000000), int64(maxChunkLimit), ...) passes, the two follow-up Greater/Less checks cannot fail and add no coverage. Consider dropping them, or replacing the exact-equality assertion with just the bounds if the intent is to tolerate future tweaks to the constant.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/string_chunking_test.go` around lines 155 - 162, The test
TestMaxChunkLimitSafety currently asserts exact equality on maxChunkLimit and
then redundantly checks bounds with require.Greater/require.Less; either remove
the redundant Greater/Less calls or change the equality assertion to only assert
the acceptable range. Edit the TestMaxChunkLimitSafety function to keep a single
style: if you want a fixed sentinel value keep require.Equal(t, int64(1000000),
int64(maxChunkLimit), ...) and delete the subsequent
require.Greater/require.Less lines, or instead replace the require.Equal line
with two range checks using require.GreaterOrEqual/require.LessOrEqual (or
require.Greater/require.Less) against the chosen lower/upper bounds to allow
future tweaks to maxChunkLimit.

dumpling/export/string_chunking.go (1)

58-111: Reuse beginChunkTracking instead of open-coding the chunk-tracking handshake.

dumpling/export/dump.go:1013-1029 exposes beginChunkTracking(meta) which does exactly what Lines 59-60 + Lines 95-110 do here (register chunkStats in d.chunkedTables, return a finalizer that sets finalized=true, atomic LoadAndDelete to claim finishedTablesCounter once sent == finished). Reusing it keeps the handshake contract in one place — otherwise any future fix to the producer-side termination protocol has to be mirrored in two call sites.

♻️ Proposed refactor

 	// Initialize chunk tracking
-	chunkStats := newTableChunkStat()
-	d.chunkedTables.Store(meta.ChunkKey(), chunkStats)
+	finishChunkTracking := d.beginChunkTracking(meta)
@@
 	defer func() {
 		// Send any remaining buffered chunk as the last chunk
 		if buffer != nil {
 			_ = sendBufferedChunk(true) // Mark as last chunk
 		}
 
-		chunkStats.finalized.Store(true)
-
 		tctx.L().Debug("streaming chunking complete",
 			zap.String("database", db),
 			zap.String("table", tbl),
 			zap.Int64("totalChunks", totalChunks))
 
-		if chunkStats.finished.Load() == chunkStats.sent.Load() {
-			// Atomic claim via LoadAndDelete: if the last writer callback
-			// already handled termination, Delete returns loaded==false
-			// and we skip the increment to avoid double-counting the
-			// finishedTablesCounter.
-			if _, loaded := d.chunkedTables.LoadAndDelete(meta.ChunkKey()); loaded {
-				IncCounter(d.metrics.finishedTablesCounter)
-			}
-		}
+		finishChunkTracking()
 	}()

As per coding guidelines: "Code SHOULD remain maintainable for future readers with basic TiDB familiarity" and follow existing package-local conventions.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/string_chunking.go` around lines 58 - 111, Replace the ad-hoc
chunk-tracking handshake (creating chunkStats, storing into d.chunkedTables, and
the deferred finalizer that sets finalized and does LoadAndDelete/IncCounter)
with the existing beginChunkTracking(meta) helper: call beginChunkTracking(meta)
to obtain the chunkStats and the returned finalizer function, keep the
bufferedChunk/sendBufferedChunk logic and totalChunks variable as-is, and in the
defer first flush any remaining buffered chunk then invoke the returned
finalizer instead of the current block that sets chunkStats.finalized and does
LoadAndDelete/IncCounter; remove the duplicated manual d.chunkedTables.Store and
finalization code so the handshake is centralized in beginChunkTracking.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

dumpling/export/string_chunking.go (1)

244-259: Optional: deduplicate the per-chunk task construction.

The two branches differ only in the WHERE-clause builder; the buildWhereCondition / buildSelectQuery / newTaskTableData calls are identical. Extracting the where-clause selection makes intent slightly clearer and removes the copy-paste.

♻️ Suggested refactor

-		// Create task for chunk using previousBoundary -> currentBoundary (with buffering)
-		var newTask *TaskTableData
-
-		if len(previousBoundary) == 0 {
-			// First chunk: everything up to first boundary
-			whereClause := buildUpperBoundWhereClause(orderByColumns, currentBoundary)
-			fullWhere := buildWhereCondition(conf, whereClause)
-			query := buildSelectQuery(db, tbl, selectField, "", fullWhere, orderByClause)
-			newTask = d.newTaskTableData(meta, newTableData(query, selectLen, false), int(totalChunks), -1)
-		} else {
-			// Intermediate chunk: between previousBoundary and currentBoundary
-			whereClause := buildBoundedWhereClause(orderByColumns, previousBoundary, currentBoundary)
-			fullWhere := buildWhereCondition(conf, whereClause)
-			query := buildSelectQuery(db, tbl, selectField, "", fullWhere, orderByClause)
-			newTask = d.newTaskTableData(meta, newTableData(query, selectLen, false), int(totalChunks), -1)
-		}
+		// Create task for chunk using previousBoundary -> currentBoundary (with buffering).
+		// First chunk: everything strictly less than the first boundary.
+		// Intermediate chunk: bounded range [previousBoundary, currentBoundary).
+		var whereClause string
+		if len(previousBoundary) == 0 {
+			whereClause = buildUpperBoundWhereClause(orderByColumns, currentBoundary)
+		} else {
+			whereClause = buildBoundedWhereClause(orderByColumns, previousBoundary, currentBoundary)
+		}
+		fullWhere := buildWhereCondition(conf, whereClause)
+		query := buildSelectQuery(db, tbl, selectField, "", fullWhere, orderByClause)
+		newTask := d.newTaskTableData(meta, newTableData(query, selectLen, false), int(totalChunks), -1)

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@dumpling/export/string_chunking.go` around lines 244 - 259, The per-chunk
task construction duplicates identical calls; refactor by computing the WHERE
clause into a single variable using buildUpperBoundWhereClause(orderByColumns,
currentBoundary) when previousBoundary is empty else
buildBoundedWhereClause(orderByColumns, previousBoundary, currentBoundary), then
call buildWhereCondition(conf, whereClause), buildSelectQuery(db, tbl,
selectField, "", fullWhere, orderByClause) and create the task with
d.newTaskTableData(meta, newTableData(query, selectLen, false),
int(totalChunks), -1) only once to remove the duplicate blocks.

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes